import os
import tvm
import time
import numpy as np
from tvm.contrib import graph_runtime

def benchmark(network, size, root):
    prefix = '{}/{}.{},{}'.format(root, network, size[0], size[1])
    json = open(prefix + '.json').read()
    lib = tvm.module.load(prefix + '.so')
    params = bytearray(open(prefix + '.params', 'rb').read())
    input_data = np.random.uniform(size=(1,3,*size)).astype(np.float32)
    module = graph_runtime.create(json, lib, tvm.cpu())
    module.load_params(params)
    t0 = time.time()
    for _ in range(100):
        module.run(data=input_data)
    print('{}: {:.3f}'.format(network, (time.time() - t0) * 10))

if __name__ == '__main__':
    root = 'models/tvm'
    size = (400, 300)
    for network in ['pnet', 'dpnet']:
        benchmark(network, size, root)
